import numpy as np
import scipy.stats as stats
import matplotlib.pyplot as plt

# Sample data for Bernoulli trials
data = np.array([1, 0, 1, 1, 0, 1, 0, 1, 1, 1])  # 1s represent success, 0s represent failure
n = len(data)  # Total number of trials
k = np.sum(data)  # Number of successes

# Prior distribution (Beta distribution parameters)
alpha_prior = 2  # Prior successes
beta_prior = 2   # Prior failures

# Posterior distribution parameters
alpha_post = alpha_prior + k
beta_post = beta_prior + (n - k)

# Posterior distribution
posterior = stats.beta(alpha_post, beta_post)

# Plot the prior and posterior distributions
x = np.linspace(0, 1, 100)
prior = stats.beta(alpha_prior, beta_prior)
plt.plot(x, prior.pdf(x), label=f'Prior Beta({alpha_prior}, {beta_prior})', color='blue')
plt.plot(x, posterior.pdf(x), label=f'Posterior Beta({alpha_post}, {beta_post})', color='red')

# Add labels and legends
plt.title('Prior and Posterior Distributions for Bernoulli Trials')
plt.xlabel('p (Probability of Success)')
plt.ylabel('Density')
plt.legend()
plt.show()

# Compute and print the posterior mean and credible interval
posterior_mean = posterior.mean()
credible_interval = posterior.interval(0.95)

print(f"Posterior mean of p: {posterior_mean}")
print(f"95% credible interval for p: {credible_interval}")
